from __future__ import annotations

import json
import subprocess
from pathlib import Path
from shlex import split
from typing import TYPE_CHECKING

from pydantic import BaseModel
from semver import Version

from analysis.plugin import AnalysisPluginV0
from helperFunctions.fileSystem import get_src_dir

if TYPE_CHECKING:
    from io import FileIO

SHELL_SCRIPT = Path(get_src_dir()) / 'bin' / 'checksec'
RESULT_TO_SUMMARY = {
    'stripped': {
        True: 'STRIPPED SYMBOLS enabled',
        False: 'STRIPPED SYMBOLS disabled',
    },
    'relro': {
        'fully enabled': 'RELRO fully enabled',
        'partially enabled': 'RELRO partially enabled',
        'disabled': 'RELRO disabled',
    },
    'pie': {
        'enabled': 'PIE enabled',
        'disabled': 'PIE disabled',
        'DSO': 'PIE/DSO present',
        'REL': 'PIE/REL present',
    },
}
RELRO_OUTPUT_TO_RESULT = {
    'full': 'fully enabled',
    'partial': 'partially enabled',
    'no': 'disabled',
}
PIE_OUTPUT_TO_RESULT = {
    'yes': 'enabled',
    'no': 'disabled',
    'dso': 'DSO',
    'rel': 'REL',
}


class AnalysisPlugin(AnalysisPluginV0):
    class Schema(BaseModel):
        canary: bool
        nx: bool
        pie: str
        relro: str
        clangcfi: bool
        safestack: bool
        stripped: bool
        runpath: bool
        rpath: bool

    def __init__(self):
        super().__init__(
            metadata=(
                self.MetaData(
                    name='exploit_mitigations',
                    description='analyses ELF binaries within a firmware for present exploit mitigation techniques',
                    version=Version(1, 0, 0),
                    mime_whitelist=[
                        'application/x-executable',
                        'application/x-object',
                        'application/x-pie-executable',
                        'application/x-sharedlib',
                    ],
                    Schema=self.Schema,
                )
            )
        )
        if not SHELL_SCRIPT.is_file():
            raise RuntimeError(f'checksec not found at path {SHELL_SCRIPT}. Please re-run the backend installation.')

    def analyze(self, file_handle: FileIO, virtual_file_path: dict, analyses: dict[str, BaseModel]) -> Schema:
        del virtual_file_path, analyses
        checksec_result = _execute_checksec_script(str(file_handle.name))
        return self.Schema(
            canary=checksec_result['canary'] == 'yes',
            nx=checksec_result['nx'] == 'yes',
            pie=PIE_OUTPUT_TO_RESULT[checksec_result['pie']],
            relro=RELRO_OUTPUT_TO_RESULT[checksec_result['relro']],
            clangcfi=checksec_result['clangcfi'] == 'yes',
            safestack=checksec_result['safestack'] == 'yes',
            stripped=checksec_result['symbols'] == 'no',
            runpath=checksec_result['runpath'] == 'yes',
            rpath=checksec_result['rpath'] == 'yes',
        )

    def summarize(self, result: Schema) -> list[str]:
        return [self._construct_summary_item(feature, value) for feature, value in result.model_dump().items()]

    @staticmethod
    def _construct_summary_item(feature: str, value: str | bool) -> str:
        if feature in RESULT_TO_SUMMARY:
            return RESULT_TO_SUMMARY[feature][value]
        return f'{feature.upper()} {"enabled" if value else "disabled"}'


def _execute_checksec_script(file_path: str) -> dict[str, str]:
    checksec_process = subprocess.run(
        split(f'{SHELL_SCRIPT} --file={file_path} --format=json --extended'),
        capture_output=True,
        text=True,
        check=True,
    )
    checksec_output = json.loads(checksec_process.stdout)
    return checksec_output[file_path]
